import io
import os
import sys
import time

import magic
import pyzipper
import requests
import yaml
from pycti import OpenCTIConnectorHelper, get_config_variable


class MalwareBazaarRecentAdditions:
    """
    Process recent additions to Malware Bazaar
    """

    def __init__(self):
        # Instantiate the connector helper from config
        config_file_path = os.path.dirname(os.path.abspath(__file__)) + "/config.yml"
        config = (
            yaml.load(open(config_file_path), Loader=yaml.FullLoader)
            if os.path.isfile(config_file_path)
            else {}
        )
        self.helper = OpenCTIConnectorHelper(config)
        self.identity = self.helper.api.identity.create(
            type="Organization",
            name="MalwareBazaar",
            description="For more info, see https://bazaar.abuse.ch/about/",
        )

        self.api_url = get_config_variable(
            "MALWAREBAZAAR_RECENT_ADDITIONS_API_URL",
            ["malwarebazaar_recent_additions", "api_url"],
            config,
        )

        self.cooldown_seconds = get_config_variable(
            "MALWAREBAZAAR_RECENT_ADDITIONS_COOLDOWN_SECONDS",
            ["malwarebazaar_recent_additions", "cooldown_seconds"],
            config,
        )
        self.cooldown_seconds = int(self.cooldown_seconds)

        self.labels_color = get_config_variable(
            "MALWAREBAZAAR_RECENT_ADDITIONS_LABELS_COLOR",
            ["malwarebazaar_recent_additions", "labels_color"],
            config,
        )

        self.include_tags = get_config_variable(
            "MALWAREBAZAAR_RECENT_ADDITIONS_INCLUDE_TAGS",
            ["malwarebazaar_recent_additions", "include_tags"],
            config,
        )
        if self.include_tags:
            self.include_tags = self.include_tags.split(",")

        self.include_reporters = get_config_variable(
            "MALWAREBAZAAR_RECENT_ADDITIONS_INCLUDE_REPORTERS",
            ["malwarebazaar_recent_additions", "include_reporters"],
            config,
        )
        if self.include_reporters:
            self.include_reporters = self.include_reporters.split(",")

        labels = get_config_variable(
            "MALWAREBAZAAR_RECENT_ADDITIONS_LABELS",
            ["malwarebazaar_recent_additions", "labels"],
            config,
        )
        self.label_ids = []

        # Create default labels
        if labels:
            labels = labels.split(",")
            for label in labels:
                created_label = self.helper.api.label.create(
                    value=label, color=self.labels_color
                )
                self.label_ids.append(created_label["id"])

    def run(self):
        self.helper.log_info("Starting MalwareBazaar Recent Additions Connector")
        while True:
            try:
                recent_additions_list = self.get_recent_additions()
                for recent_additions_dict in recent_additions_list:
                    self.helper.log_info(f"Processing: {recent_additions_dict}")
                    sha256 = recent_additions_dict["sha256_hash"]
                    reporter = recent_additions_dict["reporter"]
                    file_name = recent_additions_dict["file_name"]
                    tags = (
                        recent_additions_dict["tags"]
                        if recent_additions_dict["tags"]
                        else []
                    )

                    # If the artifact has an excluded tag, skip processing
                    if self.include_reporters:
                        if reporter not in self.include_reporters:
                            self.helper.log_info(
                                f"Skipping {sha256} as it was from a reporter not in the included list: {reporter}"
                            )
                            continue

                    if self.include_tags:
                        if not any(x in tags for x in self.include_tags):
                            self.helper.log_info(
                                f"Skipping {sha256} as it did not contain a tag in the included list."
                            )
                            continue

                    # If the artifact already exists in OpenCTI skip it
                    if self.artifact_exists_opencti(sha256):
                        self.helper.log_info(
                            f'Skipping Artifact with "{sha256}" as it already exists in OpenCTI.'
                        )
                        continue

                    # Download the artifact and unzip with default "infected" password
                    file_contents = self.download_unzip(sha256)

                    # Upload the artifact to OpenCTI
                    response = self.upload_artifact_opencti(
                        file_name,
                        file_contents,
                        f"Uploaded to MalwareBazaar by Twitter user: {reporter}.",
                    )

                    # Create external reference to MalwareBazaar report
                    external_reference = self.helper.api.external_reference.create(
                        source_name="MalwareBazaar Recent Additions",
                        url=f"https://bazaar.abuse.ch/sample/{sha256}/",
                        description="MalwareBazaar Recent Additions",
                    )
                    self.helper.api.stix_cyber_observable.add_external_reference(
                        id=response["id"],
                        external_reference_id=external_reference["id"],
                    )

                    # Attach all default labels if any
                    for label_id in self.label_ids:
                        self.helper.api.stix_cyber_observable.add_label(
                            id=response["id"], label_id=label_id
                        )

                    # Attach all tags as labels if any
                    for tag in tags:
                        label = self.helper.api.label.create(
                            value=tag,
                            color=self.labels_color,
                        )
                        self.helper.api.stix_cyber_observable.add_label(
                            id=response["id"], label_id=label["id"]
                        )

                self.helper.log_info(
                    f"Re-checking for new additions in {self.cooldown_seconds} seconds..."
                )

            except (KeyboardInterrupt, SystemExit):
                self.helper.log_info("Connector stop")
                sys.exit(0)

            except Exception as e:
                self.helper.log_error(str(e))

            if self.helper.connect_run_and_terminate:
                self.helper.log_info("Connector stop")
                sys.exit(0)

            time.sleep(self.cooldown_seconds)

    def get_recent_additions(self):
        """
        Get recent additions to MalwareBazaar.

        See https://bazaar.abuse.ch/api/#latest_additions

        returns: a dict containing the recent additions in the last
                 60 minutes to MalwareBazaar.
        """

        data = {"query": "get_recent", "selector": "time"}
        resp = requests.post(self.api_url, data=data)

        # Handle the response data

        recent_additions_list = resp.json()
        return recent_additions_list["data"]

    def download_unzip(self, sha256):
        """
        Download and unzip a sample from MalwareBazaar.

        sha256: a str representing the sample's sha256.
        returns: a bytes object containing the contents of the file
        """
        data = {"query": "get_file", "sha256_hash": sha256}
        resp = requests.post(self.api_url, data=data)
        zip_contents = resp.content
        zip_obj = io.BytesIO(zip_contents)
        zip_file = pyzipper.AESZipFile(zip_obj)
        zip_file.setpassword(b"infected")
        file_name = zip_file.namelist()[0]
        return zip_file.read(file_name)

    def artifact_exists_opencti(self, sha256):
        """
        Determine whether or not an Artifact already exists in OpenCTI.

        sha256: a str representing the sha256 of the artifact's file contents
        returns: a bool indicidating the aforementioned
        """

        custom_attributes = """
            id
            entity_type
        """
        response = self.helper.api.stix_cyber_observable.read(
            filters=[{"key": "hashes_SHA256", "values": [sha256]}],
            customAttributes=custom_attributes,
        )

        if response:
            return True
        return False

    def upload_artifact_opencti(self, file_name, file_contents, description):
        """
        Upload a file to OpenCTI.

        file_name: a str representing the name of the file
        file_contents: a bytes object representing the file contents
        description: a str representing the description for the upload

        returns: response of upload
        """

        mime_type = magic.from_buffer(file_contents, mime=True)

        kwargs = {
            "file_name": file_name,
            "data": file_contents,
            "mime_type": mime_type,
            "x_opencti_description": description,
        }

        return self.helper.api.stix_cyber_observable.upload_artifact(**kwargs)


if __name__ == "__main__":
    try:
        malwarebazaar_recent_additions = MalwareBazaarRecentAdditions()
        malwarebazaar_recent_additions.run()
    except Exception as e:
        print(e)
        time.sleep(10)
        sys.exit(0)
